from libc cimport stdlib

from bourbon.job cimport INITIALIZED, STARTED, FINISHED, FINALIZED
	
cdef class Reporter(BourbonObject):

	def __init__(self, int step):
		BourbonObject.__init__(self)
		self._step = step
	
	cpdef set_step(self, int step):
		self._step = step
	
	cpdef int step(self):
		return self._step
		
	cpdef start(self, Simulator ps):
		pass

	cpdef record(self, Simulator ps):
		pass

	cpdef finish(self, Simulator ps):
		pass
		
	cpdef object params(self):
		return {}
		
	cpdef object results(self):
		return {}
		
	def __getstate__(self):
		data = BourbonObject.__getstate__(self)
		data['step'] = self._step
		return data

	def __setstate__(self, data):
		BourbonObject.__setstate__(self, data)
		self._step = data['step']

	def __reduce__(self):
		return (Reporter,(self._step,), self.__getstate__() )
			
cdef class Simulation(Job):
	
	def __cinit__(self):
		self._creporters = NULL
		
	def __init__(self):
		Job.__init__(self)
		self._reporters = []
		
	def __dealloc__(self):
		#DBGprint 'Experiment.DEALLOC'
		if self._creporters <> NULL: stdlib.free(self._creporters)
		
	cpdef Simulator simulator(self):
		return self._simulator

	cpdef set_simulator(self, Simulator ps):
		if self.status <> INITIALIZED: raise Exception("Cannot change Simulator mid-run.")
		if self._simulator is not ps:
			ps.set_simulation(self)
		self._simulator = ps
	
	cpdef int start(self):
		if self._simulator is None: raise Exception("Cannot start experiment without simulator.")
		if not Job.start(self): return 0
		
		self._create_reporters()
		return 1
		
	cpdef object results(self):
		"""Return a dictionary of results."""
		cdef object result = Job.results(self)
		cdef Reporter r
		for r in self._reporters:
			result.update( r.results() )
		result.update( self._simulator.results() )
		
		return result
		
	
	cpdef object params(self):
		"""Return a dictionary of parameters."""
		cdef object param = Job.params(self)
		cdef Reporter r
		for r in self._reporters:
			param.update( r.params() )
			
		param.update(self._simulator.params())
		
		return param
		
	cdef _create_reporters(self):
		cdef Reporter r
		cdef int i
		self._num_reporters = len(self._reporters)
		self._creporters = <PyObject**>stdlib.malloc(sizeof(PyObject*)*self._num_reporters)
		for i in range(self._num_reporters):
			self._creporters[i] = <PyObject*>self._reporters[i]
			self._reporters[i].start(self._simulator)
	
		
	cpdef int finalize(self):
		if not Job.finalize(self): return 0
	
		cdef int i
		for i in range(self._num_reporters):
			self._reporters[i].finish(self._simulator)
		
		if self._creporters <> NULL:
			#DBGprint 'Experiment.FINALIZE'
			stdlib.free(self._creporters)
			self._creporters = NULL
		
		return 1
		
	# TODO: this is messed up
	cpdef int record(self):
		cdef int i
		cdef Reporter r
		for i in range(self._num_reporters):
			r = <Reporter>self._creporters[i]
			if r._step > 0 and self._simulator.clock() % r._step == 0:
				r.record(self._simulator)
		return 0
		
	cpdef int check(self):
		"""Returns 0 when the simulator should halt evolution, 1 otherwise."""
		return 1

	cpdef add_reporter(self, Reporter r):
		if self.status <> INITIALIZED: raise Exception("Cannot add reporter during the experiment")
		
		self._reporters.append(r)
		
		
	cpdef reporters(self):
		return self._reporters
		
	
	def __getstate__(self):
		data = Job.__getstate__(self)
		data['reporters'] = self._reporters
		data['simulator'] = self._simulator
		return data

	def __setstate__(self, data):
		Job.__setstate__(self, data)
		self._reporters = data['reporters']
		self._simulator = data['simulator']
		if self.status == STARTED:
			self._create_reporters()
		
	def __reduce__(self):
		return (Simulation,(), self.__getstate__() )
		
cdef class Simulator(BourbonObject):
	
	cdef void set_simulation(self, Simulation s):
		self._simulation = s
		
	cpdef int clock(self):
		return 0
	
	def __getstate__(self):
		return dict(simulation=self._simulation)
		
	def __setstate__(self, data):
		self._simulation = data['simulation']
		
	def __reduce__(self):
		return (Simulator, (), self.__getstate__())
		
